#!/usr/bin/python

import sys, os

if len(sys.argv) != 5:
    print "check_sge_custom -w <warning> -c <critical>"
    sys.exit(3)

W = int(sys.argv[2])
C = int(sys.argv[4])

cmd = "qstat | awk ' { print $4 } ' | sed -e \"1,2d\" | sort | uniq"
users = [u.strip() for u in os.popen(cmd).readlines()]

cmd = "qstat"
lines = [l.strip() for l in os.popen(cmd).readlines()]
counts = {}
for line in lines[2:]:
    toks = line.split()
    state = toks[4]
    user  = toks[3]
    if 'r' in state:
        state = 'running'
    elif 'q' in state:
        state = 'queued'
    else:
        state = 'unknown'
    slots = toks[-2]
    task  = toks[-1]
    task = task.split(':')
    counts[user] = counts.get(user, {})
    if len(task) <= 1:
        counts[user][state] = counts[user].get(state, 0) + 1
    else:
        start,end=task[0].split('-')
        counts[user][state] = counts[user].get(state, 0) + \
                    (int(end) - int(start) + 1)

total = sum([sum(v.values()) for k, v in counts.iteritems()])
entries = [",".join( [str(ele) for ele in 
    [k, v.get('running', 0), v.get('queued', 0), v.get('unknown', 0)]]) \
    for k, v in counts.iteritems()]

output = ":".join(entries)

if total < W:
    print "SGE_CUSTOM OK - (%i jobs)|%s" % (total, output)
    sys.exit(0)
if total < C:
    print "SGE_CUSTOM WARN - (%i jobs)|%s" % (total, output)
    sys.exit(1)

print "SGE_CUSTOM CRIT - (%i jobs)|%s" % (total, output)
sys.exit(2)

